import time
import torch

from .plugins_base import PluginBase
from trainer.trainer import PluginType, TrainContext

from models.metrics.segmentation import SegmentationMetric


class ValEvaluationPlugin(PluginBase):
    plugin_hooks = {
        PluginType.EPOCH_END: "evaluate"
    }

    def __init__(self):
        self.history = []

    def evaluate(self, ctx: TrainContext):
        if (ctx.epoch + 1) % 3 != 0:
            ctx.workspace['val_miou'] = None
            return
        
        model = ctx.model
        val_loader = ctx.val_loader
        criterion = ctx.criterion
        device = ctx.device

        if model is None or val_loader is None:
            raise RuntimeError("Validation plugin requires model and val_loader in context.")

        model.eval()

        miou_calculator = SegmentationMetric(num_classes=ctx.cfg["num_classes"], ignore_index=255)

        total_loss = 0.0

        start_time = time.time()
        with torch.no_grad():
            for inputs, targets in val_loader:
                inputs = inputs.to(device)
                targets = targets.to(device)

                outputs = model(inputs)['out']
                predictions = torch.argmax(outputs, dim=1)

                metrics = miou_calculator.update(predictions, targets)

                loss = criterion(outputs, targets)
                total_loss += loss.item()
        
        duration = time.time() - start_time

        val_loss = total_loss / len(val_loader)
        val_miou = metrics["mIoU"]
        val_mpa = metrics["mPA"]

        self.history.append({
            "epoch": ctx.epoch,
            "val_loss": val_loss,
            "val_miou": val_miou,
            "val_mpa": val_mpa
        })

        msg = {
            "mode": "val",
            "epoch": ctx.epoch + 1,
            "lr": round(ctx.optimizer.param_groups[0]['lr'], 6),
            "time": round(duration, 6),
            "loss": round(val_loss, 6),
            "acc": round(val_mpa, 6),
            "miou": round(val_miou, 6),
        
        }
        if self.check_key(ctx.workspace, "logger"):
            ctx.workspace["logger"](str(msg))
        ctx.workspace['val_miou'] = val_miou
        model.train()
